from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
import graphviz


iris = datasets.load_iris()
X = iris.data
y = iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

DTC = DecisionTreeClassifier()
DTC.fit(X_train, y_train)


export_graphviz(DTC, out_file = 'classify_tree.dot',
               feature_names = iris.feature_names, class_names = iris.target_names)
with open('classify_tree.dot') as f:
    dot_graph = f.read()
print(dot_graph)
graph = graphviz.Source(dot_graph)
graph.render('classify_tree')
